#include "stdafx.h"
#include "BaseCommand.h"
#include "AutoHandle.h"


AutoHandle _autoHandle;


BaseCommand::BaseCommand( const std::wstring& command )
	: m_command( command )
{
}

BaseCommand::~BaseCommand()
{
}

void BaseCommand::dumpBuffer( const BYTE* pBuffer, size_t cbBuffer )
{
	while ( cbBuffer > 0 )
	{
		size_t cbToPrint = min( cbBuffer, 8 );

		for ( size_t i = 0; i < cbToPrint; i++ )
		{
			wprintf( L"%02x ", pBuffer[i] );
		}
		for ( size_t i = cbToPrint; i < 8; i++ )
		{
			wprintf( L"   " );
		}
		wprintf( L"    " );
		for ( size_t i = 0; i < cbToPrint; i++ )
		{
			wprintf( L"%c", iswprint( pBuffer[i] ) ? pBuffer[i] : L'.' );
		}
		wprintf( L"\n" );

		cbBuffer -= cbToPrint;
		pBuffer += cbToPrint;
	}
}

std::wstring BaseCommand::GetCommand()
{
	return m_command;
}

ACCESS_MASK BaseCommand::GetAccessParameter( const std::map<std::wstring, std::wstring>& params, const std::wstring& name, const std::wstring& defaultValue )
{
	ACCESS_MASK am = 0;
	std::wstring desiredAccess = GetNamedParameter( params, name, defaultValue );
	while ( desiredAccess.size() >= 2 )
	{
		std::wstring accessPart = desiredAccess.substr( 0, 2 );
		if ( !_wcsnicmp( accessPart.c_str(), L"GA", 2 ) )
		{
			am |= GENERIC_ALL;
		}
		else if ( !_wcsnicmp( accessPart.c_str(), L"GR", 2 ) )
		{
			am |= GENERIC_READ;
		}
		else if ( !_wcsnicmp( accessPart.c_str(), L"GW", 2 ) )
		{
			am |= GENERIC_WRITE;
		}
		else if ( !_wcsnicmp( accessPart.c_str(), L"KA", 2 ) )
		{
			am |= KEY_ALL_ACCESS;
		}
		else if ( !_wcsnicmp( accessPart.c_str(), L"KR", 2 ) )
		{
			am |= KEY_READ;
		}
		else if ( !_wcsnicmp( accessPart.c_str(), L"KW", 2 ) )
		{
			am |= KEY_WRITE;
		}
		desiredAccess = desiredAccess.substr( 2 );
	}

	return am;
}

HANDLE BaseCommand::GetHandleParameter( const std::map<std::wstring, std::wstring>& params, const std::wstring& name )
{
	wchar_t* pEndScan;
	std::wstring handleValue = GetNamedParameter( params, name );
	if ( !_wcsicmp( handleValue.c_str(), L"AUTO" ) )
	{
		HANDLE handle = _autoHandle.GetAutoHandle();
		wprintf( L"AUTO-HANDLE: %d\n", handle );
		return handle;
	}
	else if ( !_wcsnicmp( handleValue.c_str(), L"AUTO-", wcslen(L"AUTO-") ) )
	{
		long number = wcstoul( handleValue.substr(5).c_str(), &pEndScan, 0 );
		HANDLE handle = _autoHandle.GetAutoHandle( number );
		wprintf( L"AUTO-HANDLE: %d\n", handle );
		return handle;
	}
	return (HANDLE)wcstoul( handleValue.c_str(), &pEndScan, 0 );
}

HANDLE BaseCommand::GetHandleParameter( const std::map<std::wstring, std::wstring>& params, const std::wstring& name, HANDLE defaultValue )
{
	wchar_t* pEndScan;
	std::wstring handleValue = GetNamedParameter( params, name, L"" );
	if ( handleValue.empty() )
	{
		return defaultValue;
	}
	else if ( !_wcsicmp( handleValue.c_str(), L"AUTO" ) )
	{
		HANDLE handle = _autoHandle.GetAutoHandle();
		wprintf( L"AUTO-HANDLE: %d\n", handle );
		return handle;
	}
	else if ( !_wcsnicmp( handleValue.c_str(), L"AUTO-", wcslen(L"AUTO-") ) )
	{
		long number = wcstoul( handleValue.substr(5).c_str(), &pEndScan, 0 );
		HANDLE handle = _autoHandle.GetAutoHandle( number );
		wprintf( L"AUTO-HANDLE: %d\n", handle );
		return handle;
	}
	return (HANDLE)wcstoul( handleValue.c_str(), &pEndScan, 0 );
}

std::wstring BaseCommand::GetNamedParameter( const std::map<std::wstring, std::wstring>& params, const std::wstring& name )
{
	return params.find(name)->second;
}

std::wstring BaseCommand::GetNamedParameter( const std::map<std::wstring, std::wstring>& params, const std::wstring& name, const std::wstring& defaultValue )
{
	std::map<std::wstring, std::wstring>::const_iterator it = params.find(name);
	if ( it == params.end() )
	{
		return defaultValue;
	}
	return it->second;
}

long BaseCommand::GetNamedParameterInt( const std::map<std::wstring, std::wstring>& params, const std::wstring& name )
{
	wchar_t* pEndScan;
	return wcstol( GetNamedParameter(params, name).c_str(), &pEndScan, 0 );
}

long BaseCommand::GetNamedParameterInt( const std::map<std::wstring, std::wstring>& params, const std::wstring& name, long defaultValue )
{
	std::map<std::wstring, std::wstring>::const_iterator it = params.find(name);
	if ( it == params.end() )
	{
		return defaultValue;
	}
	wchar_t* pEndScan;
	return wcstol( GetNamedParameter(params, name).c_str(), &pEndScan, 0 );
}

QuitCommand::QuitCommand()
	: BaseCommand( L"Quit" )
{
}

QuitCommand::~QuitCommand()
{
}

void QuitCommand::DisplayHelp()
{
	wprintf( L"Help: quit\n" );
}

bool QuitCommand::Execute( const std::map<std::wstring, std::wstring>& params )
{
	return true;
}

QueryObjectCommand::QueryObjectCommand()
	: BaseCommand( L"QueryObject" )
{
}

QueryObjectCommand::~QueryObjectCommand()
{
}

void QueryObjectCommand::DisplayHelp()
{
	wprintf( L"Help: queryobject -handle [handle] -class [information_class] -bufferlen [buffer_size]\n" );
}

bool QueryObjectCommand::Execute( const std::map<std::wstring, std::wstring>& params )
{
	if (
		params.size() != 3 ||
		params.find( L"handle" ) == params.end() ||
		params.find( L"class" ) == params.end() ||
		params.find( L"bufferlen" ) == params.end()
		)
	{
		DisplayHelp();
		return false;
	}

	std::wstring objectInformationClass = GetNamedParameter(params, L"class");
	NT::OBJECT_INFORMATION_CLASS ObjectInformationClass = NT::ObjectBasicInformation;
	if ( !_wcsicmp( objectInformationClass.c_str(), L"Name" ) )
	{
		ObjectInformationClass = NT::ObjectNameInformation;
	}
	else if ( !_wcsicmp( objectInformationClass.c_str(), L"Type" ) )
	{
		ObjectInformationClass = NT::ObjectTypeInformation;
	}
	else if ( !_wcsicmp( objectInformationClass.c_str(), L"All" ) )
	{
		ObjectInformationClass = NT::ObjectAllInformation;
	}
	else if ( !_wcsicmp( objectInformationClass.c_str(), L"Data" ) )
	{
		ObjectInformationClass = NT::ObjectDataInformation;
	}

	unsigned long cbBuffer = GetNamedParameterInt(params, L"bufferlen");
	BYTE* pBuffer = (BYTE*) malloc( cbBuffer );
	unsigned long cbResult;
	NtApi ntapi;
	NTSTATUS status = ntapi.NtQueryObject( GetHandleParameter(params, L"handle"), (OBJECT_INFORMATION_CLASS)ObjectInformationClass, pBuffer, cbBuffer, &cbResult );
	wprintf( L"Status = 0x%08x\n", status );
	wprintf( L"ResultLength = %d\n", cbResult );
	if ( NT_SUCCESS( status ) )
	{
		dumpBuffer( pBuffer, cbResult );
	}
	free( pBuffer );
	return false;
}
